# -*- coding: utf-8 -*-
'''
    尝试通过代理获取websocket
'''
import datetime
import os
import sys
from websocket import create_connection

import websockets
import mitmproxy.http
from mitmproxy import ctx
import mitmproxy.websocket
import json

print(sys.argv)
key = sys.argv[-1]  # 获取key


class Counter:
    def __init__(self):

        self.filename = 'robot{}.obj'.format(datetime.datetime.now().strftime('%Y-%m-%d'))
        self.robot = None
        self.flag = False

    def setRobot(self, robot):
        self.robot = robot
        self.flag = True

    # Websocket lifecycle
    def websocket_handshake(self, flow: mitmproxy.http.HTTPFlow):
        """
            Called when a client wants to establish a WebSocket connection. The
            WebSocket-specific headers can be manipulated to alter the
            handshake. The flow object is guaranteed to have a non-None request
            attribute.
        """
        ctx.log.info('---------------------------------')
        ctx.log.info('好像是建立连接了')
        '''接收到的消息类型 http
           接收到的消息内容 {'websocket': True}
        '''
        ctx.log.info('接收到的消息类型 %s' % flow.type)
        ctx.log.info('接收到的消息内容 %s' % flow.metadata)

    def websocket_start(self, flow: mitmproxy.websocket.WebSocketFlow):
        """
            A websocket connection has commenced.
        """
        ctx.log.info('~~~~~~')
        ctx.log.info('建立成功了')
        ctx.log.info('接收到的  %s' % flow.get_state())
        '''
        接收到的消息类型 websocket
        接收到的消息内容 {'websocket_handshake': 'fff06446-c5a8-4c32-86be-caa3385d4b7a'}
        '''
        ctx.log.info('接收到的消息类型 %s' % flow.type)
        ctx.log.info('接收到的消息内容 %s' % flow.metadata)
        # if flow.request.host != "fm.missevan.com":
        #     print('非猫耳',flow.request.path.startswith)
        #     return

    def websocket_message(self, flow: mitmproxy.websocket.WebSocketFlow):
        """
            Called when a WebSocket message is received from the client or
            server. The most recent message will be flow.messages[-1]. The
            message is user-modifiable. Currently there are two types of
            messages, corresponding to the BINARY and TEXT frame types.
        """
        ctx.log.info('~~~~~~')
        ctx.log.info('recive')
        ctx.log.info('接收到的 %s' % flow.messages[-1].content)
        msg = flow.messages[-1].content
        if is_json(msg):
            # 发送消息
            try:
                print('msg的类型是：',type(msg))
                send_msg(msg)
            except ConnectionRefusedError as c:
                print('机器人尚未启动')
        else:
            ctx.log.info("收到小星星")

    def websocket_error(self, flow: mitmproxy.websocket.WebSocketFlow):
        """
            A websocket connection has had an error.
        """

    def websocket_end(self, flow: mitmproxy.websocket.WebSocketFlow):
        """
            A websocket connection has ended.
        """


addons = [
    Counter()
]
'通过ws发送消息'


def send_msg(jsonObject):
    ws = create_connection("ws://127.0.0.1:8099/")
    print("Sending '%s'..." % jsonObject)
    ws.send(jsonObject)
    print("已发送")
    ws.close()


def is_json(myjson):
    try:
        json.loads(myjson)
    except ValueError:
        return False;
    return True
